# problem3_sequential.py
import numpy as np, random, math, pandas as pd, time

# ------------------ 常量和几何 ------------------
g = 9.8
missile_speed = 300.0
cloud_sink = 3.0
cloud_r = 10.0
cloud_duration = 20.0

fake_target = np.array([0.0, 0.0, 0.0])
real_target = np.array([0.0, 200.0, 0.0])   # 质点
M1 = np.array([20000.0, 0.0, 2000.0])
FY1 = np.array([17800.0, 0.0, 1800.0])

# 固定无人机航向朝向假目标，速度固定为 120 m/s （降维，利于接力优化）
alpha_fixed = math.atan2(-FY1[1], -FY1[0])
v_fixed = 120.0

missile_dir = fake_target - M1
missile_unit = missile_dir / np.linalg.norm(missile_dir)
missile_flight_time = np.linalg.norm(missile_dir) / missile_speed

# ------------------ 基本几何/物理函数 ------------------
def missile_pos(t):
    return M1 + missile_unit * missile_speed * t

def explosion_center(alpha, v, td, f):
    dir_xy = np.array([math.cos(alpha), math.sin(alpha)])
    drop = FY1 + np.array([v * dir_xy[0] * td, v * dir_xy[1] * td, 0.0])
    burst = drop + np.array([v * dir_xy[0] * f, v * dir_xy[1] * f, -0.5 * g * f * f])
    return burst

def cloud_center(alpha, v, td, f, tau):
    burst = explosion_center(alpha, v, td, f)
    z = max(0.0, burst[2] - cloud_sink * tau)
    return np.array([burst[0], burst[1], z])

def point_segment_dist(C, M, P):
    # distance from point C to segment M-P
    MP = P - M
    seg2 = MP.dot(MP)
    if seg2 <= 1e-12:
        return np.linalg.norm(C - M)
    t = np.dot(C - M, MP) / seg2
    t = max(0.0, min(1.0, t))
    proj = M + t * MP
    return np.linalg.norm(C - proj)

def blocked_at_time(alpha, v, td, f, t):
    te = td + f
    tau = t - te
    if tau < 0.0 or tau > cloud_duration:
        return False
    Mpos = missile_pos(t)
    Cpos = cloud_center(alpha, v, td, f, tau)
    d = point_segment_dist(Cpos, Mpos, real_target)
    return d <= cloud_r + 1e-9

# ------------------ 评价函数（硬判定） ------------------
def total_shield_time_from_plans(drops, bursts, dt=0.1):
    """drops, bursts: lists length 3 (seconds). 返回并集遮蔽时长和每弹单独时长+区间"""
    T_end = missile_flight_time
    times = np.arange(0.0, T_end + 1e-9, dt)
    shield_flags = np.zeros_like(times, dtype=bool)
    indiv_flags = [np.zeros_like(times, dtype=bool) for _ in range(3)]
    for idx, t in enumerate(times):
        tglob = float(t)
        Mpos = missile_pos(tglob)
        for i in range(3):
            tau = tglob - (drops[i] + bursts[i])
            if tau < 0 or tau > cloud_duration:
                continue
            if blocked_at_time(alpha_fixed, v_fixed, drops[i], bursts[i], tglob):
                indiv_flags[i][idx] = True
                shield_flags[idx] = True
    total = shield_flags.sum() * dt
    # build intervals for each individual
    indiv_intervals = []
    indiv_times = []
    for arr in indiv_flags:
        intervals = []
        in_seg = False
        start = None
        for k, val in enumerate(arr):
            tcur = times[k]
            if val and not in_seg:
                in_seg = True; start = tcur
            if not val and in_seg:
                # end at previous time + dt
                intervals.append((start, times[k-1] + dt))
                in_seg = False; start = None
        if in_seg:
            intervals.append((start, times[-1] + dt))
        indiv_intervals.append(intervals)
        indiv_times.append(sum(e - s for s, e in intervals))
    return total, indiv_times, indiv_intervals

# ------------------ 单枚弹优化（SA） ------------------
def optimize_one_grenade(fixed_drops, fixed_bursts, idx_to_opt,
                         td_bounds=(0, 60), f_bounds=(0.5, 10.0),
                         dt_eval=0.1,
                         sa_steps=1500, T0=1.0, cooling=0.999):
    """
    fixed_drops/ fixed_bursts: lists of length 3 where entries for other grenades are fixed,
      for the grenade to optimize put None.
    idx_to_opt: 0/1/2 index selected for optimization
    returns best (td, f) found and corresponding total_shield_time
    """
    # ensure initial feasible guess:
    # If there's a known good starting point for idx0 use (1.5,3.6), else mid bounds
    if idx_to_opt == 0:
        cur_td = 1.5
        cur_f = 3.6
    else:
        # place subsequent grenades after previous + 6s heuristic
        prev_td = max(0.0, fixed_drops[idx_to_opt - 1] if fixed_drops[idx_to_opt - 1] is not None else 2.0)
        cur_td = min(td_bounds[1], prev_td + 6.0)
        cur_f = 3.5
    # clamp into bounds
    cur_td = float(max(td_bounds[0], min(td_bounds[1], cur_td)))
    cur_f = float(max(f_bounds[0], min(f_bounds[1], cur_f)))

    # assemble fixed arrays
    drops = [fixed_drops[i] if fixed_drops[i] is not None else (cur_td if i == idx_to_opt else None) for i in range(3)]
    bursts = [fixed_bursts[i] if fixed_bursts[i] is not None else (cur_f if i == idx_to_opt else None) for i in range(3)]

    # If other entries None, initialize them reasonably (shouldn't happen in sequential flow)
    for i in range(3):
        if drops[i] is None:
            drops[i] = 0.0
        if bursts[i] is None:
            bursts[i] = 3.5

    # evaluate function
    best_td, best_f = cur_td, cur_f
    best_plan_drops = drops.copy(); best_plan_bursts = bursts.copy()
    best_plan_drops[idx_to_opt] = best_td
    best_plan_bursts[idx_to_opt] = best_f
    best_score, _, _ = total_shield_time_from_plans(best_plan_drops, best_plan_bursts, dt=dt_eval)

    T = T0
    cur_td, cur_f = best_td, best_f
    cur_score = best_score

    for step in range(sa_steps):
        # random neighbor: gaussian steps, but also allow occasional big jump
        if random.random() < 0.05:
            cand_td = random.uniform(*td_bounds)
            cand_f = random.uniform(*f_bounds)
        else:
            cand_td = cur_td + random.gauss(0, 1.0)   # 1s sigma
            cand_f = cur_f + random.gauss(0, 0.6)
        # enforce limits and interval constraints:
        cand_td = max(td_bounds[0], min(td_bounds[1], cand_td))
        cand_f = max(f_bounds[0], min(f_bounds[1], cand_f))

        # ensure min 1s gap to previous grenades
        cand_drops = best_plan_drops.copy(); cand_bursts = best_plan_bursts.copy()
        cand_drops[idx_to_opt] = cand_td
        cand_bursts[idx_to_opt] = cand_f

        # enforce ordering: for idx>0 ensure td >= previous_td + 1
        ok = True
        if idx_to_opt > 0:
            prev_td = cand_drops[idx_to_opt - 1]
            if cand_td < prev_td + 1.0:
                ok = False
        # also ensure td+f <= missile_flight_time
        if cand_td + cand_f > missile_flight_time - 0.01:
            ok = False

        if not ok:
            T *= cooling
            continue

        # evaluate
        score, _, _ = total_shield_time_from_plans(cand_drops, cand_bursts, dt=dt_eval)

        # acceptance
        if score > cur_score or random.random() < math.exp((score - cur_score) / max(1e-9, T)):
            cur_td, cur_f = cand_td, cand_f
            cur_score = score
            best_plan_drops[idx_to_opt] = cand_td
            best_plan_bursts[idx_to_opt] = cand_f

            if score > best_score:
                best_td, best_f = cand_td, cand_f
                best_score = score

        T *= cooling

    # final small grid refinement around best
    best_local_td, best_local_f, best_local_score = best_td, best_f, best_score
    grid_td = np.arange(max(td_bounds[0], best_td-2.0), min(td_bounds[1], best_td+2.0)+1e-9, 0.2)
    grid_f = np.arange(max(f_bounds[0], best_f-1.5), min(f_bounds[1], best_f+1.5)+1e-9, 0.2)
    for tdg in grid_td:
        for fg in grid_f:
            cand_drops = best_plan_drops.copy(); cand_bursts = best_plan_bursts.copy()
            cand_drops[idx_to_opt] = tdg; cand_bursts[idx_to_opt] = fg
            if idx_to_opt > 0:
                if cand_drops[idx_to_opt] < cand_drops[idx_to_opt-1] + 1.0:
                    continue
            if tdg + fg > missile_flight_time - 0.01:
                continue
            score, _, _ = total_shield_time_from_plans(cand_drops, cand_bursts, dt=dt_eval)
            if score > best_local_score:
                best_local_score = score
                best_local_td = tdg
                best_local_f = fg
    return best_local_td, best_local_f, best_local_score

# ------------------ 主：逐枚接力优化 ------------------
def sequential_optimize():
    # bounds for times
    td_bounds = (0.0, min(60.0, missile_flight_time - 1.0))
    f_bounds = (0.5, 10.0)

    # start with trivial fixed lists
    drops = [None, None, None]
    bursts = [None, None, None]

    # optimize grenade 1
    print("优化第1枚烟幕弹...")
    td1, f1, s1 = optimize_one_grenade(drops, bursts, 0, td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=2000, T0=1.2, cooling=0.9995)
    drops[0] = td1; bursts[0] = f1
    print(f"  弹1: td={td1:.2f}, f={f1:.2f}, 单独遮蔽(估)={s1:.2f}s")

    # optimize grenade 2 (ensure >=1s after grenade1)
    print("优化第2枚烟幕弹...")
    td2, f2, s2 = optimize_one_grenade(drops, bursts, 1, td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=1800, T0=1.0, cooling=0.9995)
    drops[1] = td2; bursts[1] = f2
    print(f"  弹2: td={td2:.2f}, f={f2:.2f}, 累计遮蔽(估)={s2:.2f}s")

    # optimize grenade 3
    print("优化第3枚烟幕弹...")
    td3, f3, s3 = optimize_one_grenade(drops, bursts, 2, td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=1600, T0=1.0, cooling=0.9995)
    drops[2] = td3; bursts[2] = f3
    print(f"  弹3: td={td3:.2f}, f={f3:.2f}, 累计遮蔽(估)={s3:.2f}s")

    # final precise evaluation (dt=0.02)
    total, indiv_times, indiv_intervals = total_shield_time_from_plans(drops, bursts, dt=0.02)
    print("\n=== 最终精确评估 ===")
    print(f"飞行方向 (deg): {math.degrees(alpha_fixed):.2f}, 速度 (m/s): {v_fixed:.1f}")
    for i in range(3):
        print(f"弹{i+1}: 投放={drops[i]:.2f}s, 引信={bursts[i]:.2f}s, 单独遮蔽={indiv_times[i]:.2f}s, 区间={indiv_intervals[i]}")
    print(f"并集总有效遮蔽时间: {total:.2f} s")

    # 保存
    df = pd.DataFrame({
        "弹编号": [1,2,3],
        "投放时间(s)": drops,
        "引信延时(s)": bursts,
        "单枚遮蔽秒数(精确dt=0.02)": indiv_times
    })
    df.to_excel("result3_sequential.xlsx", index=False)
    print("结果已保存到 result3_sequential.xlsx")
    return drops, bursts, total, indiv_times, indiv_intervals

if __name__ == "__main__":
    random.seed(42); np.random.seed(42)
    t0 = time.time()
    drops, bursts, total, indiv_times, indiv_intervals = sequential_optimize()
    print("总耗时: {:.1f}s".format(time.time() - t0))
